/******************************************************************************/
/* File: ctrlSocket.cpp                                                      */
/******************************************************************************/

#include "ctrlSocket.h"

/******************************************************************************/
/*    CtrlClient Class                                                        */
/******************************************************************************/

CtrlClient::CtrlClient(IPAddress &servAddr, int servPort, bool bUseIPv6)
          : cServAddr(servAddr), iServPort(servPort), TCPClientSocket(bUseIPv6) {

}

CtrlClient::~CtrlClient() {

}

void CtrlClient::ConnectToServer() {
    ConnectServer(cServAddr, iServPort);
}

bool CtrlClient::AddNewFlow(FlowConfig *pfc, MessageQueue *pMsgQueue) {
    if (pfc->iRole == 2)
        return true;
    
    IPAddress recvIP, sendIP;
    char *cip, *sip;
    int sport, rport;
    ThreadConf *ptc;
    ConnectionInfo *pci;
    
    bool ipv6 = (pfc->bUseIPv6 == 0) ? false : true;
    sendIP.SetHostName(pfc->sSenderIP, ipv6);
    recvIP.SetHostName(pfc->sRecverIP, ipv6);
    
    if (pfc->iProto == 1) {
        
        std::cout << "Add New TCP Flow......\n";
        
        // 新建TCP流量连接
        TCPSocket *pts;
        try {
            pts = new TCPSocket(ipv6);
            pts->Connect(recvIP, pfc->iRecverPort);
        }
        catch (SocketException se) {
            se.PrintException();
            return false;
        }
        pts->SetLocal();
        pts->SetConnected();
        
        cip = pts->LocalHostName();
        sip = pts->RemoteHostName();
        sport = pts->LocalPortNumber();
        rport = pts->RemotePortNumber();
        
        pci = new ConnectionInfo();
        pci->sSenderIP = strdup(cip);
        pci->sRecverIP = strdup(sip);
        pci->iSenderPort = sport;
        pci->iRecverPort = rport;
        pci->iProto = pfc->iProto;
        pci->iRole = pfc->iRole;
    
        // 如前面所说，我们这里将flowConfig的内容进行更改
        // 包括发送端端口号以及接收端端口号 -- 随机分配（因此需要进行网络获取并赋值）
        pfc->iSenderPort = sport;
        pfc->iRecverPort = rport;
        
        // 新建线程配置类对象，用于管理线程行为
        ptc = new ThreadConf();
        ptc->pFC = pfc;
        ptc->pMQ = pMsgQueue;
        
        // 启动TCP流量发送端线程
        TCPTrafficSender *pTcpSender = new TCPTrafficSender(pts, ptc);
        pTcpSender->Start(NULL);
        pTcpSender->Detach();
        
        DataStore::AddNewFlowThread(dynamic_cast<Thread *>(pTcpSender), pci, ptc);
    }
    else if (pfc->iProto == 2) {
        // 新建UDP流量连接
        std::cout << "Add New UDP Flow......\n";
        
        UDPConnectedSocket *pus;
        
        try {
            pus = new UDPConnectedSocket(ipv6);
            pus->Connect(recvIP, pfc->iRecverPort);
        }
        catch (SocketException se) {
            se.PrintException();
            return false;
        }
        
        cip = pus->LocalHostName();
        sip = pus->RemoteHostName();
        sport = pus->LocalPortNumber();
        rport = pus->RemotePortNumber();
        
        pci = new ConnectionInfo();
        pci->sSenderIP = strdup(cip);
        pci->sRecverIP = strdup(sip);
        pci->iSenderPort = sport;
        pci->iRecverPort = rport;
        pci->iProto = pfc->iProto;
        pci->iRole = pfc->iRole;
        
        // 如前面所说，我们这里将flowConfig的内容进行更改
        // 包括发送端端口号以及接收端端口号 -- 随机分配（因此需要进行网络获取并赋值）
        pfc->iSenderPort = sport;
        pfc->iRecverPort = rport;
        
        // 新建线程配置类对象，用于管理线程行为
        ptc = new ThreadConf();
        ptc->pFC = pfc;
        ptc->pMQ = pMsgQueue;
        
        // 启动UDP流量发送端线程
        UDPTrafficSender *pUdpSender = new UDPTrafficSender(pus, ptc);
        pUdpSender->Start(NULL);
        pUdpSender->Detach();
        
        DataStore::AddNewFlowThread(dynamic_cast<Thread *>(pUdpSender), pci, ptc);
    }
    
    return true;
}

bool CtrlClient::ModExistFlow(FlowConfig *pfc) {
    
    ConnectionInfo *pci = new ConnectionInfo();
    
    pci->sSenderIP = strdup(pfc->sSenderIP);
    pci->sRecverIP = strdup(pfc->sRecverIP);
    pci->iSenderPort = pfc->iSenderPort;
    pci->iRecverPort = pfc->iRecverPort;
    pci->iProto = pfc->iProto;
    pci->iRole = pfc->iRole;
    
    pthread_t pid = DataStore::FindFlowThread(pci);
    if(pid == -1)
        return false;
    
    ThreadConf *ptc = DataStore::FindFlowThreadConf(pid);
    if(ptc == NULL)
        return false;
    
    FlowConfig *oldFC = ptc->pFC;
    ptc->pFC = pfc;
    
    delete oldFC;
    return true;
}

bool CtrlClient::DelExistFlow(FlowConfig *pfc) {
    
    ConnectionInfo *pci = new ConnectionInfo();
    
    pci->sSenderIP = strdup(pfc->sSenderIP);
    pci->sRecverIP = strdup(pfc->sRecverIP);
    pci->iSenderPort = pfc->iSenderPort;
    pci->iRecverPort = pfc->iRecverPort;
    pci->iProto = pfc->iProto;
    pci->iRole = pfc->iRole;
    
    pthread_t pid = DataStore::FindFlowThread(pci);
    ThreadConf *ptc = DataStore::FindFlowThreadConf(pid);
    
    ptc->SetHaltFlag(1);
    
    while(ptc->GetStoppedFlag() != 1) {
        
    }
    
    DataStore::RemoveFlowThread(pci);
    delete pci;
    
    return true;
}

bool CtrlClient::HandleRcvdJson(cJSON *jsonObj, MessageQueue *pMsgQueue) {
    
    std::cout << cJSON_Print(jsonObj) << std::endl;
    
    Message *pMsg = new Message();
    pMsg->ParseFromJson(jsonObj);
    
    std::cout << "Handle recvd json......\n";
    
    // 分析消息内容
    /* 客户端主要处理两种消息：
       1. 注册成功消息
       2. 命令消息 
       -- ADD 会new一个sender线程
       -- MOD 会改变ThreadConf内容
       -- DEL 会关闭线程
    */
    switch (pMsg->iMessageType) {
        case (1): {
            // 处理request -- 核心（ADD/MOD/DEL线程）
            Request *request = pMsg->pRequest;
            switch (request->iRequestType) {
                case (1):
                    break;
                case (2): {
                    FlowRequest *ptr = request->pFlowHandle;
                    FlowConfig *pfc = ptr->pfc;
                    
                    switch (ptr->iHandleType) {
                        case (1): {
                            // ADD
                            AddNewFlow(pfc, pMsgQueue);       
                            break; 
                        }
                        case (2): {
                            // MOD
                            ModExistFlow(pfc);
                            break;
                        }
                        case (3): {
                            // DEL
                            DelExistFlow(pfc);
                            break;
                        }
                    }
                }
            }
            break;
        }
        case (2):{
            // 处理Reply -- 主要是注册/注销是否成功反馈
            Reply *reply = pMsg->pReply;
            switch (reply->iReplyType) {
                case (1): {
                    if (!reply->iSuccess)
                        std::cout << "Register traffic server failed.\n";
                    else
                        std::cout << "Register traffic server succeed.\n";
                        
                    break;
                }
                case (2): 
                    break;
            }
            break;
        }
        case (3):
            // 处理Report -- 并不需要，因为肯定收不到Report消息
            break;
        default :
            std::cout << "Unknown Message!\n";
            break;
    }
    
    //delete pMsg;
}

void CtrlClient::RunClient() {
    struct timeval retryTime;
    MessageQueue *pMsgQueue;
    CtrlSender *pCtrlSender;
    TCPTrafficServer *pTcpTrafficServer;
    UDPTrafficServer *pUdpTrafficServer;
    TrafficServerSetting *pTrafficServerSetting;
    
    while(true) {
        try {
            std::cout << "\nTry to connect to control server.\n";
            ConnectToServer();
            
            std::cout << "Connected to server successfully.\n";
            
            if(bIsConnected) {
                // 建立消息队列
                pMsgQueue = new MessageQueue();
                std::cout << "MessageQueue established." << std::endl;
                
                // 新建控制连接发送线程，同时将消息队列指针传入处理，再启动线程
                pCtrlSender = new CtrlSender(this, pMsgQueue);
                pCtrlSender->Start(NULL);
                pCtrlSender->Detach();
                std::cout << "CtrlSender started." << std::endl;
                
                pTrafficServerSetting = new TrafficServerSetting();
                pTrafficServerSetting->sHostName = strdup("0.0.0.0");
                
                // 启动TCP流量服务器线程
                pTcpTrafficServer = new TCPTrafficServer(pTrafficServerSetting, pMsgQueue);
                pTcpTrafficServer->Start(NULL);
                pTcpTrafficServer->Detach();
                std::cout << "TCP Traffic Server started." << std::endl;
                
                // 启动UDP流量服务器线程
                pUdpTrafficServer = new UDPTrafficServer(pTrafficServerSetting, pMsgQueue);
                pUdpTrafficServer->Start(NULL);
                pUdpTrafficServer->Detach();
                std::cout << "UDP Traffic Server started." << std::endl;
            }
            
            // 循环接收消息，然后再处理
            while(bIsConnected) {
                try {
                    cJSON *json = ReadJsonFromSocket();
                    HandleRcvdJson(json, pMsgQueue);
                }
                catch (SocketException se) {
                    if(se.CompareErrCode(SocketException::errNotConnected))
                        bIsConnected = false;
                    CloseSocket();
                    ResetSock();
                    // 关闭TCP和UDP流量服务器
                    pTcpTrafficServer->CloseListenSocket();
                    pUdpTrafficServer->CloseListenSocket();
                    pTcpTrafficServer->Cancel();
                    std::cout << "Shut down TCP Traffic Server...\n";
                    pUdpTrafficServer->Cancel();
                    std::cout << "Shut down UDP Traffic Server...\n";
                    
                    se.PrintException();
                }
            }
        } catch (SocketException se) {
           se.PrintException();
           std::cout << "Connect to Control Server failed.\n";
           std::cout << "Reconnect after 2 seconds.\n";
        }
        
        // 2秒后重连
        retryTime.tv_sec = 2;
        retryTime.tv_usec = 0;
        select(0, NULL, NULL, NULL, &retryTime);
        
    }
}

/******************************************************************************/
/*    CtrlServer Class                                                        */
/******************************************************************************/
CtrlServer::CtrlServer(CtrlServerSetting *pSet) 
                 : TCPServerSocket(pSet->bUseIPv6){
    fileName = pSet->sFileName;
    
    file.open(fileName, std::ifstream::in);
    if(!file.is_open()) {
        std::cout << "Config file open failed.\n";
        fileStatus = 0;
    }
    
    struct stat fileStat;
    if(stat(fileName, &fileStat) == -1 ) {
        std::cout << "Cannot read information from config file.\n";
        fileStatus = 0;
    }
    
    fileStatus = 1;
    
    lastModifiedTime = fileStat.st_mtime;
    
    IPAddress cLocalAddr;
    cLocalAddr.SetHostName(pSet->sHostName, pSet->bUseIPv6);
 
    try {
        const int iReUseAddrFlag = 1;
        SetSockOpt(SOL_SOCKET, SO_REUSEADDR, &iReUseAddrFlag, sizeof(iReUseAddrFlag));
        Bind(cLocalAddr, pSet->iLocalPort);
        SetLocal();
        Listen(pSet->iBackLog);
        std::cout << "Control Server is Running......\n";
    } catch (SocketException se) {
        se.PrintException();
        std::cout << "Cannot Start Control Server.\n";
        exit(-1);
    }
}

CtrlServer::~CtrlServer() {
    file.close();
}

void CtrlServer::SendMessage(Message *pm, TCPSocket *pts) {
    
    if(pm) {
        char *lip = pts->LocalHostName();
        int localPort = pts->LocalPortNumber();
        
        char *rip = pts->RemoteHostName();
        int remotePort = pts->RemotePortNumber();
        
        pm->sFromIP = strdup(lip);
        pm->iFromPort = localPort;
        pm->sToIP = strdup(rip);
        pm->iToPort = remotePort;
        
        cJSON *json = pm->GenToJson();
        std::cout << "Send conf msg:" << std::endl; 
        std::cout << cJSON_Print(json) << std::endl;
        int iTempSock = pts->GetSockDesc();
        SetNonblocking(iTempSock, 0);
        pts->WriteJsonToSocket(json);
        SetNonblocking(iTempSock, 1);
        
        cJSON_Delete(json);
        
        delete pm;
    }
}

bool CtrlServer::CheckFileModification() {
    time_t currentModifiedTime;
    struct stat fileStat;
    if(stat(fileName, &fileStat) == -1) {
        std::cout << "Cannot Read Information From Config File.\n";
        fileStatus = 0;
        return false;
    }
    
    currentModifiedTime = fileStat.st_mtime;
    double diffSecs = difftime(currentModifiedTime, lastModifiedTime);
    
    if(diffSecs > 0) {
        lastModifiedTime = currentModifiedTime;
        return true;
    }
    else
        return false;
}

// Parse new traffic configuration
bool CtrlServer::HandleNewTrafficConfig() {
    
    std::cout << "Config file content has been changed.\n";
    
    if(!file.is_open()) {
        std::cout << "Config file is not opend.\n";
        fileStatus = 0;
        
        // 重新打开文件
        file.open(fileName, std::ifstream::in);
    }
    
    if(!file.is_open()) {
        std::cout << "Cannot open config file.\n";
        return false;
    }
    
    fileStatus = 1;
    
    // 读文件然后分析文件内容
    char *configStr = ReadFileToString(file);
    cJSON *cj = cJSON_Parse(configStr);
    
    if(cj == NULL)
        return false;
         
    FlowRequest *tconf = new FlowRequest();
    tconf->ParseFromJson(cj);
    cJSON_Delete(cj);
    
    Request *pr = new Request();
    pr->iRequestType = 2;
    pr->pFlowHandle = tconf;
        
    Message *pm = new Message();
    pm->iMessageType = 1;
    pm->pRequest = pr;
    
    TCPSocket *pts = DataStore::SearchCtrlClientSocket(tconf->pfc->sSenderIP);
    if (pts == NULL)
        return false;
    
    SendMessage(pm, pts);
    
    return true;
}

char* CtrlServer::ReadFileToString(std::ifstream &file) {
    
    std::string result;
    std::string str;
    
    while( getline(file, str) ) {
        result += str;    
    }
    
    return strdup(result.c_str());
}

int CtrlServer::SetNonblocking(int fd, int nonblocking) {
    
    int flags, newflags;

    flags = fcntl(fd, F_GETFL, 0);
    
    if (flags < 0) {
        std::cout << "Operation fcntl error: fcntl(F_GETFL).\n";
        return -1;
    }
    
    if (nonblocking)
        newflags = flags | (int) O_NONBLOCK;
    else
        newflags = flags & ~((int) O_NONBLOCK);
    if (newflags != flags) {
	    if (fcntl(fd, F_SETFL, newflags) < 0) {
	        std::cout << "Operation fcntl error: fcntl(F_SETFL).\n";
	        return -1;
	    }
    }
    
    return 0;
}

void CtrlServer::RunServer() {    
    int result, maxfd;
    fd_set readSet, masterSet;
    struct timeval timeout;
    
    // 初始化select ReadSet 和 WriteSet
    FD_ZERO(&readSet);
    FD_ZERO(&masterSet);
    
    FD_SET(iSockDesc, &masterSet);
    maxfd = iSockDesc;
    
    while (true) {
        
        // 检查是否需要更新配置过程
        if(CheckFileModification() == true){
            HandleNewTrafficConfig();
        }
        
        // 设置超时时间
        timeout.tv_sec = 2;
        timeout.tv_usec = 0;
        
        readSet = masterSet;
        result = select(maxfd + 1, &readSet, NULL, NULL, &timeout);
        
        if(result < 0 && errno != EINTR) {
            std::cout << "Select exception: can not select for fds.\n";
            return;
        }
        
        //std::cout << "Result returned from select:" << result << "\n";
        
        if(result > 0) {
            if(FD_ISSET(iSockDesc, &readSet)) {
                TCPSocket *pts = AcceptClient();
                
                int skt = pts->GetSockDesc();
                SetNonblocking(skt, 1);
                FD_SET(skt, &masterSet);
                if(skt > maxfd)
                    maxfd = skt;
                
                DataStore::AddNewCtrlClient(pts);
            }
            
            std::list<TCPSocket *>::iterator iter = DataStore::clientList.begin();
            for(iter = DataStore::clientList.begin(); iter != DataStore::clientList.end(); ++iter) {
                int iTempSock = (*iter)->GetSockDesc();
                if(FD_ISSET( iTempSock, &readSet)) {
                    
                    int flag = 0;
                    ioctl(iTempSock, FIONREAD, &flag);
                    
                    if(flag == 0) {
                        FD_CLR(iTempSock, &masterSet);
                        close(iTempSock);
                        DataStore::RemoveCtrlClient(*iter);
                        break;
                    } else {
                        SetNonblocking(iTempSock, 0);
                        std::cout << "Read sock: " << iTempSock << std::endl;
                        std::cout << "Get json message: \n";
                        cJSON *json = (*iter)->ReadJsonFromSocket();
                        std::cout << cJSON_Print(json) << std::endl;
                        cJSON_Delete(json);
                        SetNonblocking(iTempSock, 1);
                    }
                }
            }
        }
    }
}

/******************************************************************************/
/*    CtrlSender Class                                                        */
/******************************************************************************/

CtrlSender::CtrlSender(CtrlClient *pcc, MessageQueue *pmq)
          : pCtrlClient(pcc), pMsgQueue(pmq) {
    
}

CtrlSender::~CtrlSender() {
    
}

void CtrlSender::Run() {
    
    char *fromIp = pCtrlClient->LocalHostName();
    char *toIp = pCtrlClient->RemoteHostName();
    int fromPort = pCtrlClient->LocalPortNumber(); 
    int toPort = pCtrlClient->RemotePortNumber();
    
    while(true) {
        Message *pMsg = pMsgQueue->GetMessage();
        
        if(pMsg) {
            std::cout << "Send message to control center...\n";
            
            pMsg->sFromIP = strdup(fromIp);
            pMsg->iFromPort = fromPort;
            pMsg->sToIP = strdup(toIp);
            pMsg->iToPort = toPort;
            
            cJSON *json = pMsg->GenToJson();
            std::cout << cJSON_Print(json) << std::endl;
            pCtrlClient->WriteJsonToSocket(json);
            cJSON_Delete(json);
            delete pMsg;
        }
    }
}

/******************************************************************************/
/* End File: ctrlSocket.cpp                                                   */
/******************************************************************************/